library(dplyr)
library(tidyr)
library(ComplexHeatmap)
library(circlize)

if (!file.exists("output")) {
  dir.create("output")
}

# load data
# Note: Mutational burdens and signature attributions can be found in the accompanying Supplementary Tables of the HNC manuscript

source("HNC_metadata_tidy.R")

TMB <- read.csv("Input_mutational_burden/Input_SBS_mutational_burden.csv")

sbsCOSMIC <- read.csv("Input_signature_attributions/Input_SBS_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>%
  mutate(`other SBS` = rowSums(select(., !c("SBS4", "SBS92", "SBS_I")))) %>%
  select(SBS4, SBS92, SBS_I, `other SBS`) %>%
  tibble::rownames_to_column("donor_id")

dbsCOSMIC <- read.csv("Input_signature_attributions/Input_DBS_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>%
  mutate(`other DBS` = rowSums(select(., !c("DBS2", "DBS6")))) %>%
  select(DBS2, DBS6, `other DBS`) %>%
  tibble::rownames_to_column("donor_id")

idCOSMIC <- read.csv("Input_signature_attributions/Input_ID_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>%
  mutate(`other ID` = rowSums(select(., !"ID3"))) %>%
  select(ID3, `other ID`) %>%
  tibble::rownames_to_column("donor_id")

additional_df <- list(sbsCOSMIC, dbsCOSMIC, idCOSMIC)

cosmic <- additional_df[[1]]
for (n in 2:length(additional_df)) {
  cosmic <- cosmic %>% left_join(additional_df[[n]], by = "donor_id")
}

for (df in list(cosmic, TMB)) {
  data <- data %>% left_join(df, by = "donor_id")
}

# Prepare signature matrices for heatmap

dataM <- t(as.matrix(data %>% select_at(vars(names(cosmic))) %>% tibble::column_to_rownames("donor_id")))

# Heatmap column order

col_order <- data %>%
  tibble::column_to_rownames("donor_id") %>%
  arrange(desc(tobacco), subsite, desc(Mutational.burden))

# build heatmap

colors_vir <- rev(scales::viridis_pal()(7))
SBScol <- c(colors_vir[2:4], "white")
DBScol <- c(colors_vir[5:6], "white")
IDcol <- c(colors_vir[7], "white")

hmap_annotation <- HeatmapAnnotation(
  annotation_name_rot = 0,
  "Subsite" = anno_block(labels = c("Oral cavity", "Oropharynx", "HPX", "Larynx"), labels_gp = gpar(fontsize = 10, fontface = "bold"), height = unit(0.5, "cm"), gp = gpar(col = "white")),
  "Tobacco" = data$tobacco,
  "TMB" = anno_barplot(data$Mutational.burden, baseline = 0, height = unit(1, "cm"), gp = gpar(fill = "#595959", lwd = 0), bar_width = 1, axis_param = list(at = c(0, 200000))),
  "SBS4" = anno_barplot(data$SBS4, baseline = 0, height = unit(1, "cm"), gp = gpar(fill = colors_vir[2], lwd = 0), bar_width = 1, axis_param = list(at = c(0, 60000))),
  "SBS92" = anno_barplot(data$SBS92, baseline = 0, height = unit(1, "cm"), gp = gpar(fill = colors_vir[3], lwd = 0), bar_width = 1, axis_param = list(at = c(0, 30000))),
  "SBS_I" = anno_barplot(data$SBS_I, baseline = 0, height = unit(1, "cm"), gp = gpar(fill = colors_vir[4], lwd = 0), bar_width = 1, axis_param = list(at = c(0, 15000))),
  "DBS2" = anno_barplot(data$DBS2, baseline = 0, height = unit(1, "cm"), gp = gpar(fill = colors_vir[5], lwd = 0), bar_width = 1, axis_param = list(at = c(0, 600))),
  "DBS6" = anno_barplot(data$DBS6, baseline = 0, height = unit(1, "cm"), gp = gpar(fill = colors_vir[6], lwd = 0), bar_width = 1, axis_param = list(at = c(0, 300))),
  "ID3" = anno_barplot(data$ID3, baseline = 0, height = unit(1, "cm"), gp = gpar(fill = colors_vir[7], lwd = 0), bar_width = 1, axis_param = list(at = c(0, 3000))),
  show_legend = c(
    "Subsite" = F,
    "Tobacco" = T,
    "TMB" = F,
    "SBS4" = F,
    "SBS92" = F,
    "SBS_I" = F,
    "DBS2" = F,
    "DBS6" = F,
    "ID3" = F
  ),
  col = list("Tobacco" = c("Current smoker" = "#2c7fb8", "Ex-smoker" = "#7fcdbb", "Never" = "#edf8b1")),
  gap = unit(c(1.2, 1.2, 1.2, rep(0, 2), 1.2, 0, 1.2), "mm"),
  na_col = "#e6e6e6",
  border = TRUE,
  annotation_name_side = "left",
  annotation_name_gp = gpar(fontsize = 10),
  simple_anno_size = unit(0.6, "cm")
)

HM <- Heatmap(dataM,
  # Titles
  column_title = "Mutograph HNC (n = 265)", column_title_side = "top",
  row_title_rot = 90,
  row_title_gp = gpar(fontsize = 10, fontface = "bold", col = "white"),
  # Row and column features
  show_column_names = F,
  show_row_names = F,
  column_order = rownames(col_order),
  # Heatmap size
  height = unit(0, "cm"),
  border = F,
  # Turn off clustering
  cluster_columns = F,
  cluster_rows = F,
  # Splitting
  column_split = data$subsite,
  # Legend
  show_heatmap_legend = F,
  heatmap_legend_param = list(title = "Average \n relative prop."),
  # Annotations
  bottom_annotation = hmap_annotation
)

# Legend

lgd <- Legend(
  title = "Signatures",
  labels = c("SBS4", "SBS92", "SBS_I", "DBS2", "DBS6", "ID3"),
  legend_gp = gpar(fill = colors_vir[-1])
)

myLegends <- packLegend(lgd, direction = "vertical")

# Save figure

pdf(
  file = paste0("output/", "Figure_3a_", Sys.Date(), ".pdf"),
  width = 9.5, height = 10
)
draw(HM,
  heatmap_legend_side = "right",
  annotation_legend_side = "right",
  merge_legend = TRUE,
  annotation_legend_list = myLegends
)
dev.off()
